package main

import (
    "adkr/pkg/adkr"
    "adkr/pkg/crypto"
    "bytes"
    "encoding/json"
    "flag"
    "fmt"
    "io"
    "log"
    "os"
    "strings"
    "sync"
    "sync/atomic"
    "time"
)

type bus struct {
    n int
    ins []chan adkr.Message
    outs []chan adkr.Message
    stopCh chan struct{}
    wg sync.WaitGroup
    sentCnt []int64
    recvCnt []int64
}

func newBus(n int) *bus {
    b := &bus{ n: n, ins: make([]chan adkr.Message, n), outs: make([]chan adkr.Message, n), stopCh: make(chan struct{}), sentCnt: make([]int64, n), recvCnt: make([]int64, n) }
    for i := 0; i < n; i++ { b.ins[i] = make(chan adkr.Message, 1024); b.outs[i] = make(chan adkr.Message, 1024) }
    return b
}

func (b *bus) start() {
    for i := 0; i < b.n; i++ {
        b.wg.Add(1)
        idx := i
        go func() {
            defer b.wg.Done()
            for {
                select {
                case <-b.stopCh:
                    return
                case msg := <-b.outs[idx]:
                    if msg.To >= 0 && msg.To < b.n {
                        msg.Sender = idx
                        atomic.AddInt64(&b.sentCnt[idx], 1)
                        atomic.AddInt64(&b.recvCnt[msg.To], 1)
                        b.ins[msg.To] <- msg
                    } else {
                        for j := 0; j < b.n; j++ {
                            cm := msg
                            cm.Sender = idx
                            atomic.AddInt64(&b.sentCnt[idx], 1)
                            atomic.AddInt64(&b.recvCnt[j], 1)
                            b.ins[j] <- cm
                        }
                    }
                }
            }
        }()
    }
}

func (b *bus) stop() { close(b.stopCh); b.wg.Wait() }

func (b *bus) snapshot() (sent, recv []int64) {
    sent = make([]int64, b.n); recv = make([]int64, b.n)
    for i := 0; i < b.n; i++ { sent[i] = atomic.LoadInt64(&b.sentCnt[i]); recv[i] = atomic.LoadInt64(&b.recvCnt[i]) }
    return
}

type metricsCollector struct {
    w io.Writer
    mu sync.Mutex
    coin map[int]map[string]interface{}
}

func newMetricsCollector(w io.Writer) *metricsCollector { return &metricsCollector{ w: w, coin: make(map[int]map[string]interface{}) } }

func (mc *metricsCollector) Write(p []byte) (int, error) {
    if mc.w != nil { _, _ = mc.w.Write(p) }
    line := string(bytes.TrimSpace(p))
    if strings.Contains(line, "ThresholdCoin-METRICS ") {
        idx := strings.Index(line, "ThresholdCoin-METRICS ")
        if idx >= 0 {
            js := strings.TrimSpace(line[idx+len("ThresholdCoin-METRICS "):])
            var m map[string]interface{}
            if json.Unmarshal([]byte(js), &m) == nil {
                if v, ok := m["pid"].(float64); ok { mc.mu.Lock(); mc.coin[int(v)] = m; mc.mu.Unlock() }
            }
        }
    }
    return len(p), nil
}

func main() {
    n := flag.Int("n", 4, "number of nodes")
    f := flag.Int("f", 1, "fault tolerance")
    round := flag.Int("round", 0, "coin round")
    skipVerify := flag.Bool("skip-verify", true, "skip aggregate signature verify for bench")
    flag.Parse()

    mc := newMetricsCollector(os.Stderr)
    log.SetOutput(mc)

    // crypto setup
    g1 := crypto.GetG1()
    thpkPtr, thpks, sks := crypto.GenerateThresholdKeys(*n, *f)
    thpk := thpkPtr

    // set coin skip verify to ease bench closure
    adkr.SetCoinSkipVerify(*skipVerify)

    // network bus
    b := newBus(*n); b.start(); defer b.stop()

    // launch coin protocols
    msgIns := make([]chan adkr.Message, *n)
    for i := 0; i < *n; i++ { msgIns[i] = make(chan adkr.Message, 1024) }

    for pid := 0; pid < *n; pid++ {
        pid := pid
        go func() {
            tcp := adkr.NewThresholdCoinProtocol(
                "coin-bench", pid, *n, *f, 0, nil, &g1, *round, thpks, sks[pid], thpk,
                msgIns[pid], b.outs[pid], log.Default(),
            )
            // forward messages from bus to local inbox
            go func() { for m := range b.ins[pid] { msgIns[pid] <- m } }()
            _, _ = tcp.Run(*round)
        }()
    }

    // wait until all nodes output metrics or timeout
    deadline := time.Now().Add(5 * time.Second)
    for {
        mc.mu.Lock(); cnt := len(mc.coin); mc.mu.Unlock()
        if cnt >= *n || time.Now().After(deadline) { break }
        time.Sleep(50 * time.Millisecond)
    }

    sent, recv := b.snapshot()
    summary := map[string]interface{}{
        "n": *n, "f": *f, "round": *round,
        "messages_sent": sent, "messages_recv": recv,
        "coin_metrics": mc.coin,
    }
    jb, _ := json.Marshal(summary)
    fmt.Println(string(jb))
}


